from com.demo1.ai.kimi import ai_base_config as config


## 非 Stream
# completion = client.chat.completions.create(
#     model="moonshot-v1-8k",
#     messages=[
#         {"role": "system",
#          "content": "你是 Kimi，由 Moonshot AI 提供的人工智能助手，你更擅长中文和英文的对话。你会为用户提供安全，有帮助，准确的回答。同时，你会拒绝一切涉及恐怖主义，种族歧视，黄色暴力等问题的回答。Moonshot AI 为专有名词，不可翻译成其他语言。"},
#         {"role": "user", "content": "编程判断 3214567 是否是素数。"}
#     ],
#     tools=[{
#         "type": "function",
#         "function": {
#             "name": "CodeRunner",
#             "description": "代码执行器，支持运行 python 和 javascript 代码",
#             "parameters": {
#                 "properties": {
#                     "language": {
#                         "type": "string",
#                         "enum": ["python", "javascript"]
#                     },
#                     "code": {
#                         "type": "string",
#                         "description": "代码写在这里"
#                     }
#                 },
#                 "type": "object"
#             }
#         }
#     }],
#     temperature=0.3,
# )

# print(completion.choices[0].message)


def no_stream_v1(text_content):
    completion = config.client.chat.completions.create(
        model="moonshot-v1-8k",
        messages=[
            {
                "role": "system",
                "content": "你是 Kimi，由 Moonshot AI 提供的人工智能助手，你更擅长中文对话。你会为用户提供安全，有帮助，准确的回答。同时，你会拒绝一切涉及恐怖主义，种族歧视，黄色暴力等问题的回答。你需要以英文回复我。",
            },
            {"role": "user", "content": text_content}
        ],
        temperature=0.3,
    )
    # print(completion.choices[0].message.content)
    return completion.choices[0].message.content


## stream
def stream_v1(text_content):
    print("stream_v1 start")
    response = config.client.chat.completions.create(
        model="moonshot-v1-8k",
        messages=[
            {
                "role": "system",
                "content": "你是 Kimi，由 Moonshot AI 提供的人工智能助手，你更擅长中文对话。你会为用户提供安全，有帮助，准确的回答。同时，你会拒绝一切涉及恐怖主义，种族歧视，黄色暴力等问题的回答。Moonshot AI 为专有名词，不可翻译成其他语言。",
            },
            {"role": "user", "content": text_content},
        ],
        temperature=0.3,
        stream=True,
    )
    return response


def stream_msg_print():
    response = stream_v1("你好，我叫李雷，1+1等于多少？")
    collected_messages = []
    for idx, chunk in enumerate(response):
        # print("Chunk received, value: ", chunk)
        chunk_message = chunk.choices[0].delta
        if not chunk_message.content:
            continue
        collected_messages.append(chunk_message)  # save the message
        # print(f"#{idx}: {''.join([m.content for m in collected_messages])}")
        # print("#{}: {}".format(idx, chunk_message.content))
    # print(f"Full conversation received: {''.join([m.content for m in collected_messages])}")


# stream_msg_print()